import copy

import numpy as np

from src.formation import Formation
from src.model import Model
from src.parameter import Parameter
from ..orientation_controller import OrientationController
from ..velocity_controller import VelocityController


class FormationController:
    def __init__(self, kps: list, kis: list, kds: list, model: Model, count: int, ts=0.001, position_gain=0.0) -> None:
        super().__init__()
        self._model = copy.deepcopy(model)
        self._count = count
        self._velocity_controllers = [
            VelocityController(kps[0: 3], kis[0: 3], kds[0: 3], self._model, ts=ts, position_gain=position_gain) for _
            in range(self._count)]
        self._orientation_controllers = [OrientationController(kps[3: 6], kis[3: 6], kds[3: 6], self._model, ts=ts) for
                                         _ in range(self._count)]

    def control(self, parameters: list[Parameter], formation: Formation):

        for i in range(1, self._count):
            parameters[i].dposd = parameters[0].dposd
            parameters[i].psid = parameters[0].psid

        position = [np.array([parameters[i].pos - parameters[j].pos for j in range(self._count)]).T for i in
                    range(self._count)]
        deltas = formation.cal_deltas(parameters[0].psi)

        e = [np.sum(deltas[i] - position[i], 1) for i in range(self._count)]

        out = np.zeros(self._count * 4)

        for i in range(self._count):
            u1 = self._velocity_controllers[i].control(parameters[i], e[i])
            orientation_controller_out = self._orientation_controllers[i].control2(parameters[i])
            control_out = np.array([u1, *orientation_controller_out])
            torques = self._model.assign(control_out)
            out[4 * i: 4 * (i + 1)] = torques
        return out
